from tqdm import tqdm
from utils.config import *

# 训练
def train(model, train_loader, criterion, optimizer, step, writer):
    for batch in tqdm(train_loader): # 创建进度条
        optimizer.zero_grad() # 梯度清零

        # 从batch中取样本
        left_img = batch['left'].to(device)
        right_img = batch['right'].to(device)
        target_disp = batch['disp'].to(device)
        # img_name = batch['name'][0]

        # 掩码，过滤视差图中<0的值
        mask = (target_disp < max_disp) & (target_disp > 0) # 返回布尔张量，>0True
        mask = mask.detach_()

        # 前向传播
        disp = model(left_img, right_img)

        # 计算损失
        loss = criterion(disp[mask], target_disp[mask]) # 只对mask=True的值位置计算损失

        # 反向传播
        loss.backward()
        optimizer.step()

        # 绘制live_loss图像
        writer.add_scalar("live_loss", loss.item(), step)

        step += 1

    return step
